package io.piper.datasource;

import com.mysql.cj.jdbc.MysqlDataSource;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.reflect.MethodSignature;
import org.mybatis.spring.SqlSessionFactoryBean;
import org.mybatis.spring.mapper.MapperScannerConfigurer;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.ClassPathResource;
import org.springframework.jdbc.datasource.DataSourceTransactionManager;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.stereotype.Component;

import javax.sql.DataSource;
import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;
import java.util.HashMap;

@Configuration
public class Config {

    @Bean
    public DataSource master() {
        MysqlDataSource source = new MysqlDataSource();
        source.setURL("jdbc:mysql://localhost:3306/ds1");
        source.setUser("root");
        source.setPassword("123456");
        return source;
    }

    @Bean
    public DataSource slave() {
        MysqlDataSource source = new MysqlDataSource();
        source.setURL("jdbc:mysql://localhost:3306/ds2");
        source.setUser("root");
        source.setPassword("123456");
        return source;
    }

    @Bean
    public DynamicDataSource dynamicDataSource() {
        DynamicDataSource dataSource = new DynamicDataSource();
        dataSource.setDefaultTargetDataSource(master());
        HashMap<Object, Object> map = new HashMap<>();
        map.put(DataSourceEnum.MASTER, master());
        map.put(DataSourceEnum.SLAVE, slave());
        dataSource.setTargetDataSources(map);
        dataSource.afterPropertiesSet();
        return dataSource;
    }

    public static class DynamicDataSource extends AbstractRoutingDataSource {
        @Override
        protected Object determineCurrentLookupKey() {
            return DataSourceTypeManager.get();
        }
    }

    @Bean
    public SqlSessionFactoryBean sqlSessionFactory() {
        SqlSessionFactoryBean bean = new SqlSessionFactoryBean();
        bean.setDataSource(this.dynamicDataSource());
        bean.setMapperLocations(new ClassPathResource("io/piper/server/spring/pojo/mapper/xml/UserMapper.xml"));
        return bean;
    }

    @Bean
    public MapperScannerConfigurer mapperScannerConfigurer() {
        MapperScannerConfigurer config = new MapperScannerConfigurer();
        config.setBasePackage("io/piper/server/spring/pojo/mapper");
        return config;
    }

    @Bean
    public DataSourceTransactionManager transactionManager() {
        DataSourceTransactionManager manager = new DataSourceTransactionManager();
        manager.setDataSource(master());
        return manager;
    }

    public static class DataSourceTypeManager {
        private static final ThreadLocal<DataSourceEnum> dataSourceTypes = ThreadLocal.withInitial(() -> DataSourceEnum.MASTER);

        public static DataSourceEnum get() {
            return dataSourceTypes.get();
        }

        public static void set(DataSourceEnum dataSourceType) {
            dataSourceTypes.set(dataSourceType);
        }

        public static void reset() {
            dataSourceTypes.set(DataSourceEnum.MASTER);
        }
    }

    public static enum DataSourceEnum {
        MASTER, SLAVE
    }

    @Target(ElementType.METHOD)
    @Retention(RetentionPolicy.RUNTIME)
    public static @interface Routing {
        DataSourceEnum value() default DataSourceEnum.MASTER;
    }

    @Component
    @org.aspectj.lang.annotation.Aspect
    public static class Aspect {
        @Around("@annotation(io.piper.datasource.Config.Routing)")
        public void before(ProceedingJoinPoint jp) {
            MethodSignature signature = (MethodSignature) jp.getSignature();
            Routing routing = signature.getMethod().getAnnotation(Routing.class);
            try {
                DataSourceEnum before = DataSourceTypeManager.get();
                DataSourceTypeManager.set(routing.value());
                jp.proceed();
                DataSourceTypeManager.set(before);
            } catch (Throwable throwable) {
                throwable.printStackTrace();
            }
        }
    }
}
